package com.cch.cooperation.api.common.sdk.openai;

import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;


/**
 * open ai 客户端
 *
 * @author cch
 */
public class OpenAiClient {

    private String apiUrl;
    private String apiKey;

    public OpenAiClient(String apiUrl, String apiKey) {
        this.apiUrl = apiUrl;
        this.apiKey = apiKey;
    }

    public String infer(String model, String text) {
        JSONObject param = new JSONObject();
        param.put("model", model);
        param.put("messages", Lists.newArrayList(ImmutableMap.of("role", "user", "content", text)));
        param.put("temperature", 0.0);
        param.put("top_p", 0.9);
        param.put("top_k", 2);
        param.put("seed", 30);
        HttpRequest request = HttpRequest.post(apiUrl);
        request.header("Authorization", "Bearer " + apiKey);
        request.header("Content-Type", "application/json");
        request.body(param.toJSONString());
        try (HttpResponse resp = request.execute()) {
            JSONObject json = JSON.parseObject(resp.body());
            return json.getJSONArray("choices")
                    .getJSONObject(0)
                    .getJSONObject("message")
                    .getString("content");
        }

    }
}
